from langchain.prompts.example_selector import SemanticSimilarityExampleSelector
from langchain_community.vectorstores import Chroma
from langchain_huggingface import HuggingFaceEmbeddings

class find_example:
    """
    该类利用嵌入模型从示例中提取例子，处理为可供prompt导入的数据
    """
    def __init__(self, **kwargs):
        self.embedding = HuggingFaceEmbeddings(
            model_name="../../all_miniLM_L6",  # 本地模型文件夹路径
            model_kwargs={'device': 'cpu'}
        )
        self.example_selector = SemanticSimilarityExampleSelector.from_examples(
            kwargs['examples'],  # 示例库（待筛选的候选示例列表）
            self.embedding,  # 嵌入模型（用于将文本转为向量）
            Chroma,  # 向量存储（用于存储和搜索向量）
            k=1  # 筛选数量（返回最相似的1个示例）
        )
        pass
    def get_example(self,question)->str:
        """
        根据问题的向量在示例中寻找向量相似度最高的例子
        :param question: 要输入的问题
        :return: 相似度最高的例子
        """
        selected_examples = self.example_selector.select_examples({"question": question})
        return selected_examples